package com.hczt.xhminiapp.adminapi.conf.security;

import com.hczt.xhminiapp.common.bean.CodeMsg;
import com.hczt.xhminiapp.common.bean.RtnResult;
import org.springframework.http.HttpHeaders;
import org.springframework.http.ResponseEntity;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.provider.error.AbstractOAuth2SecurityExceptionHandler;
import org.springframework.security.web.AuthenticationEntryPoint;
import org.springframework.util.StringUtils;

import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;

/**
 * @author 红创-马海强
 * @date 2019-04-01 13:01
 * @说明：增强response
 */
public class OAuthExceptionEntryPoint extends AbstractOAuth2SecurityExceptionHandler implements AuthenticationEntryPoint {

    private String typeName = OAuth2AccessToken.BEARER_TYPE;

    private String realmName = "oauth";

    @Override
    public void commence(HttpServletRequest request, HttpServletResponse response, AuthenticationException authException) throws IOException, ServletException {
        doHandle(request, response, authException);
    }

    @Override
    protected ResponseEntity<?> enhanceResponse(ResponseEntity<?> response, Exception authException) {
        HttpHeaders headers = response.getHeaders();
        String existing = null;
        if (headers.containsKey("WWW-Authenticate")) {
            existing = extractTypePrefix(headers.getFirst("WWW-Authenticate"));
        }
        StringBuilder builder = new StringBuilder();
        builder.append(typeName+" ");
        builder.append("realm=\"" + realmName + "\"");
        if (existing!=null) {
            builder.append(", "+existing);
        }
        HttpHeaders update = new HttpHeaders();
        update.putAll(response.getHeaders());
        update.set("WWW-Authenticate", builder.toString());
        CodeMsg codeMsg;
        switch (response.getStatusCode()) {
            case UNAUTHORIZED:
                codeMsg = CodeMsg.OAUTH2_INVALID_TOKEN.fillArgs(authException.getMessage());
                break;
            case FORBIDDEN:
                codeMsg = CodeMsg.OAUTH2_ACCESS_DENIED.fillArgs(authException.getMessage());
                break;
            default:
                codeMsg = CodeMsg.OAUTH2_UNKNOWN_ERROR.fillArgs(authException.getMessage());
        }
        return new ResponseEntity<>(RtnResult.error(codeMsg), update, response.getStatusCode());
    }

    private String extractTypePrefix(String header) {
        String existing = header;
        String[] tokens = existing.split(" +");
        if (tokens.length > 1 && !tokens[0].endsWith(",")) {
            existing = StringUtils.arrayToDelimitedString(tokens, " ").substring(existing.indexOf(" ") + 1);
        }
        return existing;
    }
}
